import pandas as pd
from sqlalchemy import create_engine, text
import tkinter as tk
from tkinter import ttk, messagebox
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib import rcParams

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 数据库配置
DB_CONFIG = {
    'host': '127.0.0.1',
    'port': 3306,
    'user': 'root',
    'password': '',
    'database': 'stocks',
    'charset': 'utf8mb4'
}


def fetch_stocks_info(engine):
    """获取stocks_info表所有数据，并按symbol列排序"""
    query = text("SELECT * FROM stocks_info ORDER BY symbol")
    df = pd.read_sql(query, engine)
    # 添加序号列，从1开始
    df.insert(0, '序号', range(1, len(df) + 1))
    return df


def fetch_daily_data(engine, ts_code):
    """获取指定股票的日线数据，并按交易日降序排序"""
    query = text("SELECT * FROM stocks_daily WHERE ts_code = :ts_code ORDER BY trade_date DESC")
    return pd.read_sql(query, engine, params={'ts_code': ts_code})


def display_daily_data(parent, engine, ts_code, stock_name):
    """显示指定股票的日线数据"""
    try:
        data = fetch_daily_data(engine, ts_code)

        if data.empty:
            messagebox.showinfo("提示", f"没有找到 {ts_code} 的日线数据")
            return

        # 创建新窗口
        daily_window = tk.Toplevel(parent)
        daily_window.title(f"{stock_name} ({ts_code}) - 日线数据 (共 {len(data)} 条)")
        daily_window.geometry("1300x800")

        # 主框架
        main_frame = ttk.Frame(daily_window)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        # 创建Treeview组件
        tree = ttk.Treeview(main_frame, show="headings")
        tree["columns"] = list(data.columns)

        # 设置列格式
        col_widths = {
            'ts_code': 100,
            'trade_date': 100,
            'open': 80,
            'high': 80,
            'low': 80,
            'close': 80,
            'pre_close': 80,
            'chg_amo': 80,
            'pct_chg': 80,
            'vol': 100,
            'amount': 120
        }

        for col in data.columns:
            width = col_widths.get(col, 100)
            # 修改此处，使数据居中对齐
            tree.column(col, width=width, anchor='center')
            tree.heading(col, text=col, anchor='center')

        # 添加数据
        for i, row in data.iterrows():
            tree.insert("", tk.END, values=list(row))

        # 滚动条
        vsb = ttk.Scrollbar(main_frame, orient="vertical", command=tree.yview)
        hsb = ttk.Scrollbar(main_frame, orient="horizontal", command=tree.xview)
        tree.configure(yscrollcommand=vsb.set, xscrollcommand=hsb.set)

        # 布局
        tree.grid(row=0, column=0, sticky='nsew')
        vsb.grid(row=0, column=1, sticky='ns')
        hsb.grid(row=1, column=0, sticky='ew')

        # 网格配置
        main_frame.grid_rowconfigure(0, weight=1)
        main_frame.grid_columnconfigure(0, weight=1)

        # 添加图表
        if len(data) > 0:
            try:
                fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), gridspec_kw={'height_ratios': [3, 1]})

                # 价格走势图 - 反转数据顺序，使最新数据在右侧
                reversed_data = data.iloc[::-1]
                numeric_cols = ['open', 'high', 'low', 'close']
                for col in numeric_cols:
                    if col in reversed_data.columns:
                        ax1.plot(reversed_data['trade_date'], reversed_data[col], label=col)

                ax1.set_title(f"{stock_name} ({ts_code}) - 价格走势")
                ax1.legend()
                ax1.grid(True)
                plt.setp(ax1.get_xticklabels(), rotation=45)

                # 成交量柱状图 - 反转数据顺序
                if 'vol' in reversed_data.columns:
                    ax2.bar(reversed_data['trade_date'], reversed_data['vol'], width=0.6, color='blue', alpha=0.6)
                    ax2.set_title("成交量")
                    ax2.grid(True)
                    plt.setp(ax2.get_xticklabels(), rotation=45)

                plt.tight_layout()

                # 嵌入图表
                chart_frame = ttk.Frame(daily_window)
                chart_frame.pack(fill=tk.BOTH, expand=True)
                canvas = FigureCanvasTkAgg(fig, master=chart_frame)
                canvas.draw()
                canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

            except Exception as e:
                print(f"创建图表时出错: {e}")

    except Exception as e:
        messagebox.showerror("错误", f"获取日线数据时出错:\n{str(e)}")


def display_stocks_info(engine):
    """显示stocks_info表数据"""
    try:
        info_data = fetch_stocks_info(engine)

        if info_data.empty:
            messagebox.showinfo("提示", "stocks_info表中没有数据")
            return

        # 创建主窗口
        root = tk.Tk()
        root.title("股票基本信息 (双击查看日线数据)")
        root.geometry("1200x800")

        # 主框架
        main_frame = ttk.Frame(root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        # 创建Treeview组件
        tree = ttk.Treeview(main_frame, show="headings")
        tree["columns"] = list(info_data.columns)

        # 设置列格式
        col_widths = {
            '序号': 60,  # 新增的序号列
            'ts_code': 80,
            'symbol': 80,
            'name': 80,
            'area': 80,
            'industry': 100,
            'cnspell': 100,
            'market': 60,
            'list_date': 100,
            'st_status': 70,
            'is_hs': 40,
            'act_name': 200,
            'act_ent_type': 80,
            'exchange': 80
        }

        for col in info_data.columns:
            width = col_widths.get(col, 120)
            # 修改此处，使数据居中对齐
            tree.column(col, width=width, anchor='center')
            tree.heading(col, text=col, anchor='center')

        # 添加数据
        for i, row in info_data.iterrows():
            tree.insert("", tk.END, values=list(row))

        # 双击事件
        def on_double_click(event):
            item = tree.selection()[0]
            values = tree.item(item, 'values')
            ts_code = values[1]  # ts_code现在是第二列（因为第一列是序号）
            stock_name = values[3] if len(values) > 3 else ts_code  # name现在是第四列
            display_daily_data(root, engine, ts_code, stock_name)

        tree.bind("<Double-1>", on_double_click)

        # 滚动条
        vsb = ttk.Scrollbar(main_frame, orient="vertical", command=tree.yview)
        hsb = ttk.Scrollbar(main_frame, orient="horizontal", command=tree.xview)
        tree.configure(yscrollcommand=vsb.set, xscrollcommand=hsb.set)

        # 布局
        tree.grid(row=0, column=0, sticky='nsew')
        vsb.grid(row=0, column=1, sticky='ns')
        hsb.grid(row=1, column=0, sticky='ew')

        # 网格配置
        main_frame.grid_rowconfigure(0, weight=1)
        main_frame.grid_columnconfigure(0, weight=1)

        # 添加搜索功能
        search_frame = ttk.Frame(root)
        search_frame.pack(fill=tk.X, padx=10, pady=5)

        ttk.Label(search_frame, text="搜索:").pack(side=tk.LEFT, padx=5)
        search_var = tk.StringVar()
        search_entry = ttk.Entry(search_frame, textvariable=search_var)
        search_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)

        def update_tree():
            search_term = search_var.get().lower()
            for item in tree.get_children():
                values = tree.item(item, 'values')
                if any(search_term in str(v).lower() for v in values):
                    tree.item(item, tags=('match',))
                    tree.selection_set(item)
                else:
                    tree.item(item, tags=('no_match',))

            tree.tag_configure('match', background='yellow')
            tree.tag_configure('no_match', background='white')

        search_var.trace_add("write", lambda *args: update_tree())

        try:
            root.mainloop()
        except KeyboardInterrupt:
            print("程序被手动中断。")
            root.destroy()

    except Exception as e:
        messagebox.showerror("错误", f"程序初始化失败:\n{str(e)}")


if __name__ == "__main__":
    try:
        # 创建数据库连接
        engine = create_engine(
            f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@"
            f"{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}?"
            f"charset={DB_CONFIG['charset']}"
        )

        # 测试连接
        with engine.connect() as conn:
            conn.execute(text("SELECT 1"))

        # 显示股票信息
        display_stocks_info(engine)

    except Exception as e:
        messagebox.showerror("错误", f"数据库连接失败:\n{str(e)}")
